import cv2
import numpy as np
from config import cfg
from color_detect import Color_detect
from cifa_nn import Net
from cifar_trans import Transform
from cifardetect_color import detect
import torch
from torchvision import transforms,datasets
def main():
    cap = cv2.VideoCapture('rtsp://admin:hf62843295@192.168.0.6/Streaming/Channels/101')
    cv2.namedWindow('camera', cv2.WINDOW_NORMAL)
    classes = ('red', 'blue', 'yellow', 'other')
    while True:
        image = cap.read()[1]
        color = classes[0]
        if options == 'tranditional':
            cnts = C_detect.Color_Pretreat(image, color)

            for i, cnt in enumerate(cnts):
                scale_w_h, area, boundarea = C_detect.Cnt_Pretreat(cnt)
                if 2000 < area < 80000 and 1 < scale_w_h < 2:
                    if float(area / boundarea) > 0.5:
                        if color == 'red':
                            C_detect.R_detect(image, cnts, i)
                        elif color == 'blue':
                            C_detect.B_detect(image, cnts, i)
                        else:
                            C_detect.Yellow_detect(image, cnts, i)
        else:
            boxes, cls = det.Color_detect(image, 'red')
            for box in boxes:
                image1 = image[box[1] + 2:box[3] - 2, box[0] + 2:box[2] - 2]
                T_images = trans.trans(image1)
                predicted = trans.predict(T_images)
                if classes[predicted[0]] == color:
                    cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 225, 65), 1, cv2.LINE_AA)
        cv2.imshow('color_detect', image)
        cv2.waitKey(1)
        k = cv2.waitKey(1)
        if k == 32:
            cv2.destroyAllWindows()
            cap.release()

if __name__ == '__main__':
    options = input('tranditional or neural net')
    C_detect = Color_detect(cfg)
    det = detect(cfg)
    trans = Transform()
    trans.Load()
    main()